import oneflow as flow
import oneflow.nn as nn
from .conv import ConvFrontEnd
from oasr.frontend.base import BaseFrontEnd

# n_front_ends: 2
# frontend_params:
#     front0_type: raw
#     front0:
#         params1: a
#         params2: b
#     front1:
#         params1: c
#         params2: d


class StackMultiFrontEnd(BaseFrontEnd):
    def __init__(self, n_front_ends, frontend_params):
        super(StackMultiFrontEnd, self).__init__()

        self.ends = nn.ModuleList(
            [
                FrontEnd['front%d_type' % i](**frontend_params['front%d' % i]) for i in range(n_front_ends)
            ]
        )

    def forward(self, inputs, inputs_mask):

        for _, frontend in enumerate(self.ends):
            inputs, inputs_mask = frontend(inputs, inputs_mask)

        return inputs, inputs_mask

    def inference(self, inputs, inputs_mask, cache=None):
        raise NotImplementedError


FrontEnd = {
    'conv': ConvFrontEnd,
    'stack': StackMultiFrontEnd
}